import base64
import uvicorn
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import JSONResponse
import shutil
import os
from PIL import Image
import io
from fastapi.middleware.cors import CORSMiddleware
from search import search

app = FastAPI()

# 添加 CORS 中间件
origins = [
    "*"  # 放行所有
    # "http://localhost",
    # "http://localhost:7000",  # 前端服务
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/")
def read_root():
    return {"Hello": "World"}

@app.post("/uploadfile/")
async def upload_file(file: UploadFile = File(...)):
    # 存储图片
    with open(os.path.join("temp", file.filename), "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    # 处理图片
    try:
        image = Image.open(os.path.join("temp", file.filename))

        # 查询后将结果保存
        search(image, file.filename)

        res = Image.open(os.path.join("res", file.filename.rsplit('.', 1)[0] + "_res.png"))
        rgb_im = res.convert('RGB')

        data = io.BytesIO()
        rgb_im.save(data, "JPEG")
        data64 = base64.b64encode(data.getvalue())
        return JSONResponse(content={"image": f"data:image/jpeg;base64,{data64.decode('utf-8')}"})

    except Exception as e:
        return JSONResponse(content={"error": str(e)})

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8000, log_level="info")